import copy
import math

import common_modules

from DrawingAPIs.line_drawer import *
from DrawingAPIs.bar_drawer import *

def LayerData(layer_file):
    layer_time = {}
    x_ticklabels = []
    legends = []
    with open(layer_file, newline='') as csvfile:
        rows = csv.reader(csvfile, delimiter=',')
        for row in rows:
            if row[1] not in layer_time.keys():
                layer_time.setdefault(row[1], {})
                legends.append(row[1])
            layer_time[row[1]][row[3]] = float(row[4][:-3])
            if row[3] not in x_ticklabels:
                x_ticklabels.append(int(row[3]))
    return layer_time, x_ticklabels, legends

def DrawSpeedUpFigs(net, dev_layer_time, out_dir):
    # dev_layer_time: {dev: {layer: {#streams, exec_time}}}
    devices = list(dev_layer_time.keys())

    baseline = 0.0
    for dev in dev_layer_time.keys():
        conv_speedup_set = []
        stream_idx_set = []
        legends = list(dev_layer_time[dev].keys())
        y_max = 0.0
        x_max = 0
        for layer in dev_layer_time[dev].keys():
            conv_speedup = []
            x_values = list(dev_layer_time[dev][layer].keys())
            x_values = list(map(int, x_values))
            stream_idx_set.append(x_values)

            for stream_num in dev_layer_time[dev][layer].keys():
                if stream_num == '1':
                    baseline = dev_layer_time[dev][layer][stream_num]
                    conv_speedup.append(baseline / baseline)
                else:
                    assert(baseline > 0.0)
                    conv_speedup.append(baseline / dev_layer_time[dev][layer][stream_num])

            conv_speedup_set.append(conv_speedup)
            x_max = int(max(x_values))
            if y_max < max(conv_speedup):
                y_max = max(conv_speedup)

        if y_max > (math.ceil(y_max) - 0.5):
            y_max = math.ceil(y_max)
        else:
            y_max = math.ceil(y_max) - 0.5

        x_tick_pos = list(range(0, x_max + 1, 2))
        x_ticklabels = []
        for tick_pos in x_tick_pos:
            x_ticklabels.append(str(tick_pos))

        DrawLineFig(stream_idx_set, x_tick_pos, x_ticklabels, conv_speedup_set, legends,
                    "#Streams", "Speedup", 0, x_max + 1, 0.0, y_max, out_dir,
                    'mot_speedup_' + net + '_' + dev, allow_legend=True)

def DrawOptStreamFigs(net, dev_layer_time, out_dir):
    # dev_layer_time: {dev: {layer: {#streams, exec_time}}}
    # Sort dict by value: {k:v for k,v in sorted(dic.items(), key=lambda item: item[1])}
    x_ticklabels = list(dev_layer_time.keys())
    legends = list(dev_layer_time['K40c'].keys())
    num_dev = len(x_ticklabels)
    num_layers = len(legends)
    bar_width = 0.4
    x_interval = bar_width * (num_layers + 1)
    x_ticks = np.arange(x_interval - (num_layers / 2) * bar_width,
                        num_dev * x_interval, x_interval)

    baseline = 0.0
    conv_opt_streams = []
    y_max = 0
    for layer in dev_layer_time['K40c'].keys():
        layer_opt_streams = []
        for dev in dev_layer_time.keys():
            layer_opt_streams.append(min(dev_layer_time[dev][layer],
                                         key=dev_layer_time[dev][layer].get))

        layer_opt_streams = list(map(int, layer_opt_streams))
        if y_max < int(max(layer_opt_streams)):
            y_max = int(max(layer_opt_streams))

        conv_opt_streams.append(layer_opt_streams)

    y_max = math.ceil(y_max / 6) * 6
    DrawBarFig(x_ticks, x_ticklabels, conv_opt_streams, legends, "", "#Streams",
               bar_width, 0.0, max(x_ticks) + (num_layers / 2 + 0.7) * bar_width, 0, y_max,
               out_dir, "mot_" + net + "_OptStreams", allow_legend=True)

def DrawMotFigs(args, devices):
    net = args.mfolder[0]
    root_path = args.root[0]
    layer_time_file = net + 'LayerForwardTime-PROF.csv'

    x_ticklabels = []
    legends = []
    dev_layer_time = {}
    for dev in devices:
        dev_layer_time[dev], temp_x_ticklabels, temp_legends = LayerData(root_path+'/'+net+'/'+dev+'/'+layer_time_file)
        if not x_ticklabels:
            x_ticklabels = temp_x_ticklabels
        else:
            assert(x_ticklabels == temp_x_ticklabels)

        if not legends:
            legends = temp_legends
        else:
            assert(legends == temp_legends)

    DrawSpeedUpFigs(net, dev_layer_time, args.output[0] + '/MotFigs')
    DrawOptStreamFigs(net, dev_layer_time, args.output[0] + '/MotFigs')
